package main

import (
    "flag"
    "fmt"
    "log"
    "net"
    "time"
    
    "github.com/yangyongzhen/mqtt-broker/internal/protocol/mqtt311"
    "github.com/yangyongzhen/mqtt-broker/internal/protocol/common"
    "github.com/yangyongzhen/mqtt-broker/pkg/mqtt"
)

func main() {
    var (
        address  = flag.String("addr", "localhost:1883", "MQTT broker address")
        clientID = flag.String("client", "test-client", "Client ID")
        topic    = flag.String("topic", "test/topic", "Topic to publish/subscribe")
        message  = flag.String("msg", "Hello MQTT!", "Message to publish")
        mode     = flag.String("mode", "pub", "Mode: pub or sub")
    )
    flag.Parse()
    
    conn, err := net.Dial("tcp", *address)
    if err != nil {
        log.Fatalf("Failed to connect: %v", err)
    }
    defer conn.Close()
    
    if *mode == "pub" {
        publisher(conn, *clientID, *topic, *message)
    } else {
        subscriber(conn, *clientID, *topic)
    }
}

func publisher(conn net.Conn, clientID, topic, message string) {
    // 发送CONNECT
    connectPacket := &mqtt311.ConnectPacket{
        ProtocolName:    "MQTT",
        ProtocolVersion: 4,
        CleanSession:    true,
        KeepAlive:       60,
        ClientID:        clientID,
    }
    connectPacket.MessageType = common.CONNECT
    
    data, err := connectPacket.Encode()
    if err != nil {
        log.Fatalf("Encode CONNECT error: %v", err)
    }
    
    conn.Write(data)
    
    // 读取CONNACK
    reader := mqtt.NewPacketReader(conn)
    packet, err := reader.ReadPacket()
    if err != nil {
        log.Fatalf("Read CONNACK error: %v", err)
    }
    
    if connack, ok := packet.(*mqtt311.ConnackPacket); ok {
        if connack.ReturnCode != 0 {
            log.Fatalf("Connection rejected: %d", connack.ReturnCode)
        }
        fmt.Println("Connected successfully")
    }
    
    // 发送PUBLISH
    publishPacket := &mqtt311.PublishPacket{
        TopicName: topic,
        Payload:   []byte(message),
    }
    publishPacket.MessageType = common.PUBLISH
    publishPacket.QoS = common.QoS0
    
    data, err = publishPacket.Encode()
    if err != nil {
        log.Fatalf("Encode PUBLISH error: %v", err)
    }
    
    conn.Write(data)
    fmt.Printf("Published message to %s: %s\n", topic, message)
    
    time.Sleep(1 * time.Second)
}

func subscriber(conn net.Conn, clientID, topic string) {
    // 发送CONNECT
    connectPacket := &mqtt311.ConnectPacket{
        ProtocolName:    "MQTT",
        ProtocolVersion: 4,
        CleanSession:    true,
        KeepAlive:       60,
        ClientID:        clientID,
    }
    connectPacket.MessageType = common.CONNECT
    
    data, err := connectPacket.Encode()
    if err != nil {
        log.Fatalf("Encode CONNECT error: %v", err)
    }
    
    conn.Write(data)
    
    // 读取CONNACK
    reader := mqtt.NewPacketReader(conn)
    packet, err := reader.ReadPacket()
    if err != nil {
        log.Fatalf("Read CONNACK error: %v", err)
    }
    
    if connack, ok := packet.(*mqtt311.ConnackPacket); ok {
        if connack.ReturnCode != 0 {
            log.Fatalf("Connection rejected: %d", connack.ReturnCode)
        }
        fmt.Println("Connected successfully")
    }
    
    // 发送SUBSCRIBE
    subscribePacket := &mqtt311.SubscribePacket{
        PacketID: 1,
        Subscriptions: []common.Subscription{
            {Topic: topic, QoS: common.QoS0},
        },
    }
    subscribePacket.MessageType = common.SUBSCRIBE
    
    data, err = subscribePacket.Encode()
    if err != nil {
        log.Fatalf("Encode SUBSCRIBE error: %v", err)
    }
    
    conn.Write(data)
    fmt.Printf("Subscribed to %s\n", topic)
    
    // 持续读取消息
    for {
        packet, err := reader.ReadPacket()
        if err != nil {
            log.Printf("Read packet error: %v", err)
            break
        }
        
        if publish, ok := packet.(*mqtt311.PublishPacket); ok {
            fmt.Printf("Received message on %s: %s\n", publish.TopicName, string(publish.Payload))
        }
    }
}
